# multiplication of a matrix and vector
from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
# Initialize the vector length 
LENGTH = 1000

# Generate a matrix and a vector in root process
if rank >= 0:
    m = np.linspace(1, size * LENGTH, size * LENGTH).reshape((size, LENGTH))
    v = np.linspace(1, LENGTH, LENGTH)
else:
    m = None
    v = None 
# print(f"process {rank} matrix and vector: \n{m} \n{v}")

# initialize local m 
local_m = np.zeros(LENGTH)

# scatter m to local_m
comm.Scatter(m, local_m, root=0)
# Boardcast v to all processes
comm.Bcast(v, root=0)

# Compute multiplication of local_m and v
dot = np.dot(local_m, v)

# Gather all dot result
dotGathered = np.zeros(size)
comm.Gather(dot, dotGathered, root=0)

# print(f"Process {rank} local matrix: \n{local_m}")
# print(f"Process {rank} local matrix dot vector: \n{np.dot(local_m,v)}")

# print the result
if comm.rank == 0:
    print(f"The multiplication of matrix and vector in the parallel compute:\n{dotGathered}\n")